#!/usr/bin/python3
# -*- coding: utf-8 -*-

# @file base.py
# @brief
# @author QRS
# @version 1.0
# @date 2022-02-18 13:56

import zmq
# from omegaconf import OmegaConf
import json
from multiprocessing import Process
from frepai.constants import (
    ZMQ_SOCKET_PORT
)


class TaskBase(Process):

    zmq_ctx = zmq.Context()

    def __init__(self, logger, handler, stype, topic):
        super(TaskBase, self).__init__(name=topic)
        self.stype = stype
        self.topic = topic
        self.logger = logger
        self.handler = handler

    def init(self):
        self.logger.info(f'Task {self.topic} init')

    def quit(self):
        pass

    def run(self):
        try:
            self.init()
            sub = self.zmq_ctx.socket(zmq.SUB)
            sub.connect('tcp://{}:{}'.format('0.0.0.0', ZMQ_SOCKET_PORT))
            sub.subscribe(self.topic)
            while True:
                cfg = ''.join(sub.recv_string().split(' ')[1:])
                try:
                    self.step_run(json.loads(cfg), self.handler)
                except Exception as err:
                    self.logger.error(f'{err}')
            self.quit()
        except Exception as err:
            self.logger.error(f'{err}')
